from typing import Any, Dict

from dotenv import load_dotenv
from langchain_deepseek import ChatDeepSeek

from .BaseLLM import BaseLLM

# 加载环境变量
load_dotenv()


class DeepSeekThink(BaseLLM):
    """
    DeepSeek Think模型实现
    专门用于复杂推理和思考任务
    """

    def __init__(self, **kwargs):
        super().__init__(model_name="deepseek-reasoner", **kwargs)
        self.api_base = kwargs.get('api_base', 'https://api.deepseek.com')
        self.temperature = kwargs.get('temperature', 0.1)  # 推理任务使用较低温度
        self.max_tokens = kwargs.get('max_tokens', 8192)  # 推理任务需要更多token

    def create_llm(self, streaming: bool = False, callbacks=None, **kwargs):
        """
        创建DeepSeek Think LLM实例
        
        Args:
            streaming: 是否启用流式输出
            callbacks: 回调函数列表
            **kwargs: 其他配置参数
            
        Returns:
            ChatDeepSeek: DeepSeek Think LLM实例
        """
        # 基础配置
        config = {
            'model': self.model_name,
            'base_url': self.api_base,
            'max_retries': kwargs.get('max_retries', 1),
            'streaming': streaming,
            'temperature': kwargs.get('temperature', self.temperature),
            'max_tokens': kwargs.get('max_tokens', self.max_tokens),
        }

        if callbacks:
            config['callbacks'] = callbacks

        # 合并其他配置
        other_kwargs = {k: v for k, v in kwargs.items()
                        if k not in ['max_retries', 'temperature', 'max_tokens']}
        config.update(other_kwargs)

        llm = ChatDeepSeek(**config)
        return llm

    def get_model_info(self) -> Dict[str, Any]:
        """
        获取DeepSeek Think模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        return {
            'model_name': self.model_name,
            'provider': 'DeepSeek',
            'type': 'reasoning',
            'api_base': self.api_base,
            'temperature': self.temperature,
            'max_tokens': self.max_tokens,
            'supports_streaming': True,
            'supports_async': True,
            'description': 'DeepSeek Think是专门用于复杂推理和思考的模型，具有强大的逻辑推理能力'
        }

    def think(self, problem: str, **kwargs) -> str:
        """
        专门用于推理思考的方法
        
        Args:
            problem: 需要思考的问题
            **kwargs: 其他参数
            
        Returns:
            str: 推理结果
        """
        # 构造推理提示
        thinking_prompt = f"""
请仔细思考以下问题，展示你的推理过程：

问题：{problem}

请按照以下步骤进行思考：
1. 理解问题的核心
2. 分析相关因素
3. 逐步推理
4. 得出结论

思考过程：
"""

        return self.invoke(thinking_prompt, **kwargs)

    async def athink(self, problem: str, **kwargs) -> str:
        """
        异步推理思考方法
        
        Args:
            problem: 需要思考的问题
            **kwargs: 其他参数
            
        Returns:
            str: 推理结果
        """
        # 构造推理提示
        thinking_prompt = f"""
请仔细思考以下问题，展示你的推理过程：

问题：{problem}

请按照以下步骤进行思考：
1. 理解问题的核心
2. 分析相关因素
3. 逐步推理
4. 得出结论

思考过程：
"""

        return await self.ainvoke(thinking_prompt, **kwargs)

    def stream_think(self, problem: str, **kwargs):
        """
        流式推理思考方法
        
        Args:
            problem: 需要思考的问题
            **kwargs: 其他参数
            
        Yields:
            str: 推理过程片段
        """
        thinking_prompt = f"""
请仔细思考以下问题，展示你的推理过程：

问题：{problem}

请按照以下步骤进行思考：
1. 理解问题的核心
2. 分析相关因素
3. 逐步推理
4. 得出结论

思考过程：
"""

        for chunk in self.stream(thinking_prompt, **kwargs):
            yield chunk
